#include <iostream>
#include "BFS.hpp"
#include "CommonGPU.h"
#include <set>
#include <string>
#include <fstream>
#include <time.h>
#include <stdio.h>
#include <thrust/scan.h>
#include <thrust/remove.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/unique.h>
#include <thrust/sort.h>
#include <cuda_runtime_api.h>
#include <cuda.h>
#include <cuda_profiler_api.h>

using namespace std;

// global variables
thrust::device_vector<int>* adj_list; /* adjacency list */
thrust::device_vector<int>* adj_list_idx;  /* end index of adjacency list */
int N; /* number of nodes */
int M; /* number of edges */


BFS_GPU::BFS_GPU(int n){
	this->test_dvec = new thrust::device_vector<int>(n,100);
}

void BFS_GPU::print(int i){
	cout << "GPU class printing: " << (*test_dvec)[i] << endl;
}



void clean(){
	delete adj_list;
	delete adj_list_idx;
}

void BFS_init(string output_file){
	ifstream input(output_file.c_str());
	if(!input.is_open()){
		cout << "file not found" << endl;
		exit(1);
	}

	input >> N;
	input >> M;

	thrust::host_vector<int> h_adj_list(M);
	thrust::host_vector<int> h_adj_list_idx(N);

	for(int i = 0; i < N; i++){
		input >> h_adj_list_idx[i];
	}

	for(int i = 0; i < M; i++){
		input >> h_adj_list[i];
	}
	
	input.close();
	cout << "finishing reading file, N:" << N << " M:" << M << endl;
	adj_list = new thrust::device_vector<int>(h_adj_list.begin(),h_adj_list.end());
	adj_list_idx = new thrust::device_vector<int>(h_adj_list_idx.begin(),h_adj_list_idx.end());
}

// initialize adjacency list
void BFS_init(int n, int m, string output_file){
	ifstream input(output_file.c_str());
	if(input.is_open()){
		cout << "file existed" << endl;
		BFS_init(output_file);
		return;
	}

	N = n;
	M = m;
	vector<set<int> > temp_adj(N);
	
	// random generate edges
	srand(1);
	for(int i = 0; i < M; i++){
		bool finish = false;
		while(!finish){
			int node_x = rand() % N;
			int node_y = rand() % N;
			if(temp_adj[node_x].find(node_y) == temp_adj[node_x].end() && node_x != node_y){
				finish = true;
				temp_adj[node_x].insert(node_y);
			}
		}
		if(i % 500000 == 0) cout << "generating in iter: " << i << endl;
	}

	thrust::host_vector<int> h_adj_list;
	thrust::host_vector<int> h_adj_list_idx;
	int count = 0;

	for(int i = 0; i < N; i++){
		int adj_size = temp_adj[i].size();
		for(set<int>::iterator it = temp_adj[i].begin(); it != temp_adj[i].end(); it++){
			h_adj_list.push_back(*it);
		}
		count += adj_size;
		h_adj_list_idx.push_back(count);
		if(i % 10000 == 0) cout << "outputing adjacent index: " << i << endl;
	}

	if(count != M) cout << "something wrong with random graph generation" << endl;

	adj_list = new thrust::device_vector<int>(h_adj_list.begin(),h_adj_list.end());
	adj_list_idx = new thrust::device_vector<int>(h_adj_list_idx.begin(),h_adj_list_idx.end());

	if(!output_file.empty()){
		cout << "writing to file_output" << endl;
		ofstream output(output_file.c_str());
		output << N << endl;
		output << M << endl;
		for(int i = 0; i < N; i++){
			output << h_adj_list_idx[i] << " ";
		}
		output << endl;

		for(int i = 0; i < h_adj_list.size(); i++){
			output << h_adj_list[i] << " ";
		}
		output << endl;
		output.close();
	}
}

// print matrix
void print_adj_matrix(){
	thrust::host_vector<int> h_adj_list(adj_list->begin(),adj_list->end());
	thrust::host_vector<int> h_adj_list_idx(adj_list_idx->begin(),adj_list_idx->end());

	for(int i = 0; i < N; i++){
		int start = (i == 0) ? 0 : h_adj_list_idx[i-1];
		int end = h_adj_list_idx[i];
		cout << "node: " << i << ":";
		for(int j = start; j < end; j++){
			cout << h_adj_list[j] << " ";
		}
		cout << endl;
	}
}

void print_neighbours(int node){
	thrust::host_vector<int> h_adj_list(adj_list->begin(),adj_list->end());
	thrust::host_vector<int> h_adj_list_idx(adj_list_idx->begin(),adj_list_idx->end());

	int start = (node == 0) ? 0 : h_adj_list_idx[node-1];
	int end =  h_adj_list_idx[node];
	cout << "node: " << node << ":";
	for(int j = start; j < end; j++){
		cout << h_adj_list[j] << " ";
	}
	cout << endl;
}

// count number of new frontier
__global__ void count_frontier(int* adj_list_idx, int* frontier,int frontier_size, int* frontier_count){
	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	if(tid < frontier_size){
		int node = frontier[tid];
		int adj_start = (node == 0) ? 0 : adj_list_idx[node -1];
		int adj_end = adj_list_idx[node];
		frontier_count[tid] = adj_end - adj_start;
	}
}

__global__ void output_frontier(int* adj_list, int* adj_list_idx, int* frontier, int* frontier_idx, int* frontier_output){
	int bid = blockIdx.x;
	int tid = threadIdx.x;

	int node = frontier[bid];
	int adj_start = (node == 0) ? 0 : adj_list_idx[node -1];
	int adj_end = adj_list_idx[node];
	int size = adj_end - adj_start;
	int round = (size / blockDim.x) + (size % blockDim.x != 0);

	int output_start = (bid == 0) ? 0 : frontier_idx[bid - 1];

	//if(tid == 0){
	//	printf("frontier is %d\n",node);
	//	printf("adj_start is %d\n", adj_start);
	//	printf("adj_end is %d\n",adj_end);
	//	printf("output_start %d\n",output_start);
	//}

	for(int i = 0; i < round; i++){
		int rel_idx = i * blockDim.x + tid;
		int read_idx = adj_start + rel_idx;
		if(read_idx < adj_end){
			frontier_output[output_start + rel_idx] = adj_list[read_idx];
		}
	}
}

__global__ void check_visited(int* new_frontier, int new_frontier_size, int* visited, int* output){
	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	if(tid < new_frontier_size){
		output[tid] = visited[new_frontier[tid]] == 1;
	}
}

__global__ void mark_visited(int* new_frontier, int new_frontier_size,int* visited){
	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	if(tid < new_frontier_size){
		visited[new_frontier[tid]] = 1;
	}

}


void BFS(int root){
	int THREAD_PER_BLK = 1024;
	int BLKS = 0;

	thrust::device_vector<int> visited(N,0);
	visited[root] = 1;
	
	// setup frontier structure
	thrust::device_vector<int>* frontier = new thrust::device_vector<int>(1,root);
	int frontier_size = 1;

	cout << "start to compute" << endl;	
	vector<clock_t> per_time(5);
	clock_t first_start = clock();
	int num_rounds = 0;

	while(frontier_size > 0){
		num_rounds++;
		/*cudaDeviceSynchronize();
		cout << "visited nodes: ";
		for(int i= 0; i < N; i++){
			if(visited[i] == 1) cout << i << " ";
		}
		cout << endl;
		cout << "current frontier: ";
		for(int i = 0; i < frontier.size(); i++){
			cout << frontier[i] << " ";
		}
		cout << endl;
		getchar();*/
		clock_t start = clock();
		thrust::device_vector<int> frontier_count(frontier_size,0);
		BLKS = (frontier_size / THREAD_PER_BLK) + (frontier_size % THREAD_PER_BLK != 0);

		cudaDeviceSynchronize();
		cout << "counting frontiers" << endl;
		// count for each old frontier, the number of adjacent nodes
		count_frontier<<<BLKS,THREAD_PER_BLK>>>(
			thrust::raw_pointer_cast(adj_list_idx->data()),
			thrust::raw_pointer_cast(frontier->data()),
			frontier_size,
			thrust::raw_pointer_cast(frontier_count.data()));

		thrust::inclusive_scan(frontier_count.begin(),frontier_count.end(),frontier_count.begin());

		
		int temp_size = frontier_count[frontier_count.size()-1];
		cout << "temp size:" << temp_size << endl;
		thrust::device_vector<int> frontier_output(temp_size);
		thrust::device_vector<int> remove_indicator(temp_size);

		cout << "output frontiers" << endl;
		output_frontier<<<frontier_size,THREAD_PER_BLK>>>(
			thrust::raw_pointer_cast(adj_list->data()),
			thrust::raw_pointer_cast(adj_list_idx->data()),
			thrust::raw_pointer_cast(frontier->data()),
			thrust::raw_pointer_cast(frontier_count.data()),
			thrust::raw_pointer_cast(frontier_output.data()));

		
		// find visited frontier
		BLKS = (temp_size / THREAD_PER_BLK) + (temp_size % THREAD_PER_BLK != 0);


		check_visited<<<BLKS,THREAD_PER_BLK>>>(
			thrust::raw_pointer_cast(frontier_output.data()),
			temp_size,
			thrust::raw_pointer_cast(visited.data()),
			thrust::raw_pointer_cast(remove_indicator.data()));
		
		/*cout << "frointier temp: ";
		for(int i = 0; i < temp_size; i++){
			cout << frontier_output[i] << "(" << remove_indicator[i] << ") ";
		}
		cout << endl;*/

		//cout << "remove visited frontiers" << endl;
		// remove visited frontier
		thrust::device_vector<int>::iterator last_elem = thrust::remove_if(frontier_output.begin(),frontier_output.end(),remove_indicator.begin(),thrust::identity<int>());

		per_time[0] += clock()-start;

		/*cout << "after remove :";
		for(int i = 0; i < (last_elem-frontier_output.begin()); i++){
			cout << frontier_output[i] << " ";
		}
		cout << endl;*/
		/*cout << "before duplicate elimination: " << (last_elem - frontier_output.begin()) << endl;*/

		start = clock();

		//if(num_rounds == 3){
		//	last_elem = thrust::unique(frontier_output.begin(),last_elem);
		//}
		//// remove duplicate frointier 
		//thrust::sort(frontier_output.begin(),last_elem);
		//per_time[1] += clock()-start;


		//start = clock();
		//last_elem = thrust::unique(frontier_output.begin(),last_elem);
		//per_time[2] += clock()-start;
		//frontier_size = last_elem - frontier_output.begin();
		if(num_rounds == 3){
			last_elem = thrust::unique(frontier_output.begin(),last_elem);
			frontier_size = last_elem - frontier_output.begin();
		}
		frontier_size = remove_duplicate(frontier_output,(last_elem-frontier_output.begin()),num_rounds);
		
		start = clock();
		// update the frontier size and array
		delete frontier;
		frontier = new thrust::device_vector<int>(frontier_output.begin(),frontier_output.begin()+frontier_size);
		per_time[3] += clock()-start;

		/*int number_visited = thrust::reduce(visited.begin(),visited.end());
		cout << "new frontier_size: "<< frontier_size << " with number of visited: " << number_visited << endl;*/
		

		//cout << "update frontiers" << endl;
		cout << "number of frontiers now: "<< frontier_size << endl;
		// mark the new frontier as visisted nodes
		BLKS = (frontier_size / THREAD_PER_BLK) + (frontier_size % THREAD_PER_BLK != 0);
		mark_visited<<<BLKS,THREAD_PER_BLK>>>(
			thrust::raw_pointer_cast(frontier_output.data()),
			frontier_size,
			thrust::raw_pointer_cast(visited.data()));
	}

	for(int i = 0; i < 4; i++){
		cout << "part " << i << " takes: " << (double)per_time[i] / CLOCKS_PER_SEC << endl;
	}

	cout << "time taken: " << ((double)clock()-first_start)/CLOCKS_PER_SEC << endl;
	int number_visited = thrust::reduce(visited.begin(),visited.end());
	cout << "number of visited nodes: " << number_visited << endl;



}